import dashscope
from http import HTTPStatus
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath
from dashscope import ImageSynthesis
import requests


class Tongyi(object):
    def __init__(self, api_key):
        self.api_key = api_key

    def text_send(self, prompt, history=None):
        """文生文"""
        if history is None:
            history = []

        messages = [
            *history,
            {'role': 'user', 'content': prompt},
        ]
        # print('MESSAGE',messages)
        response = dashscope.Generation.call(
            model=dashscope.Generation.Models.qwen_turbo,
            messages=messages,
            result_format='message',  # 将返回结果格式设置为 message
            api_key=self.api_key,
        )
        # print('RESONE',response)
        if response.status_code == HTTPStatus.OK:
            return response.output.choices[0].message['content']
        else:
            return None

    def img_send(self, prompt, num=1, size='720*1280'):
        response = ImageSynthesis.call(
            model=ImageSynthesis.Models.wanx_v1,  # 模型名称
            prompt=prompt,  # 提示语
            n=num,  # 图片数量
            size=size,  # 图片尺寸
            api_key=self.api_key,
        )
        if response.status_code == HTTPStatus.OK:
            data = {}
            for result in response.output.results:
                data[PurePosixPath(unquote(urlparse(result.url).path)).parts[-1]] = result.url
            return data
        else:
            return {}


if __name__ == '__main__':

    ty = Tongyi(api_key='sk-8d774cf840004f028f6c6847c1e57b3d')
    response = ty.text_send("你好")
    print(response)

    # 文生图
    response = ty.img_send('画一张原神里八重神子的图')
    for file_name, url in response.items():
        with open('./%s' % file_name, 'wb+') as f:
            f.write(requests.get(url).content)